#!/usr/bin/env python3
"""
比较采取不同模型的性能
"""


result='''
sentence-transformers/all-MiniLM-L6-v2     cross-encoder/mmarco-mMiniLMv2-L12-H384-v1
模型              初始化时间(秒)        嵌入时间(秒)         查询时间(秒)         总时间(秒)         
=================================================================
HuggingFace     2.4385        16.6622        76.3656        95.4662
FlagModel       1.5285        27.8624        60.1269        89.5178
=================================================================

BAAI/bge-base-zh-v1.5    BAAI/bge-reranker-large
模型              初始化时间(秒)        嵌入时间(秒)         查询时间(秒)         总时间(秒)         
=================================================================
HuggingFace     2.0164        47.3763        118.1165        167.5092
FlagModel       2.0552        27.9038        60.1585        90.1175
=================================================================
'''
print(result)


from operator import rshift
import os
import sys
import time
import json
import traceback
import logging
from pathlib import Path

# 设置日志 - 同时输出到控制台和文件
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("logs/model_performance_test.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

logger.debug("脚本开始执行")

# 添加当前目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
logger.debug(f"当前目录: {current_dir}")
sys.path.append(current_dir)
logger.debug(f"Python路径: {sys.path}")

from config.config import load_config
from embedding import EmbeddingService
from langchain.schema import Document

# 确保中文显示正常
try:
    import matplotlib.pyplot as plt
    plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    logger.debug("matplotlib配置完成")
except ImportError as e:
    logger.error(f"导入matplotlib失败: {str(e)}")
    logger.error(f"异常堆栈: {traceback.format_exc()}")

# 测试数据路径
TEST_DOCS_DIR = os.path.join(current_dir, 'test_docs_large')
TEST_QUERIES_FILE = os.path.join(current_dir, 'test_queries_large.json')


# 加载测试文档
def load_test_documents(max_docs=None):
    """从目录加载测试文档"""
    docs = []
    if not os.path.exists(TEST_DOCS_DIR):
        logger.error(f"测试文档目录不存在: {TEST_DOCS_DIR}")
        return docs

    logger.info(f"从 {TEST_DOCS_DIR} 加载测试文档...")
    for i, filename in enumerate(sorted(os.listdir(TEST_DOCS_DIR))):
        if max_docs and i >= max_docs:
            break
        if filename.endswith('.md'):
            file_path = os.path.join(TEST_DOCS_DIR, filename)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                # 只取前1000个字符进行测试，避免处理过大的文档
                content = content[:1000]
                docs.append(Document(page_content=content))
                if (i+1) % 50 == 0 or (i+1) == len(os.listdir(TEST_DOCS_DIR)):
                    logger.info(f"已加载 {i+1} 个文档")
            except Exception as e:
                logger.error(f"读取文档 {file_path} 失败: {str(e)}")
    logger.info(f"共加载 {len(docs)} 个测试文档")
    return docs

# 加载测试查询
def load_test_queries(max_queries=None):
    """从文件加载测试查询"""
    queries = []
    if not os.path.exists(TEST_QUERIES_FILE):
        logger.error(f"测试查询文件不存在: {TEST_QUERIES_FILE}")
        return queries

    logger.info(f"从 {TEST_QUERIES_FILE} 加载测试查询...")
    try:
        with open(TEST_QUERIES_FILE, 'r', encoding='utf-8') as f:
            data = json.load(f)
            for i, item in enumerate(data):
                if max_queries and i >= max_queries:
                    break
                queries.append(item['query'])
        logger.info(f"共加载 {len(queries)} 个测试查询")
    except Exception as e:
        logger.error(f"读取查询文件失败: {str(e)}")
    return queries

class ModelPerformanceTester:
    def __init__(self):
        self.results = {}

    def test_config(self, config_path, config_name):
        """测试特定配置的性能"""
        # 设置环境变量以使用指定的配置文件
        os.environ['CONFIG_PATH'] = config_path
        logger.info(f"开始测试 {config_name} 配置...")
        logger.info(f"使用配置文件: {config_path}")

        # 加载配置以获取collection名称
        config = load_config(config_path=config_path)
        collection_name = config['milvus']['collection_name']
        logger.info(f"当前配置的collection名称: {collection_name}")

        # 如果是测试collection，则删除它（不计入测试时间）
        if collection_name == 'test':
            logger.info(f"发现测试collection '{collection_name}'，开始删除...")
            try:
                # 直接使用pymilvus删除集合
                from pymilvus import connections, utility
                
                # 连接Milvus
                connections.connect(
                    alias="default",
                    host=config['milvus']['connection_args']['host'],
                    port=config['milvus']['connection_args']['port']
                )
                logger.info("成功连接到Milvus服务")

                # 检查集合是否存在，明确指定使用'default'连接
                if utility.has_collection(collection_name, using="default"):
                    logger.info(f"集合 {collection_name} 存在")
                    delete_start_time = time.time()
                    # 删除集合，明确指定使用'default'连接
                    utility.drop_collection(collection_name, using="default")
                    delete_time = time.time() - delete_start_time
                    logger.info(f"成功删除测试collection '{collection_name}'，耗时: {delete_time:.4f} 秒")
                else:
                    logger.info(f"测试collection '{collection_name}' 不存在，无需删除")
                
                # 断开连接
                connections.disconnect(alias="default")
            except Exception as e:
                logger.error(f"删除测试collection '{collection_name}' 失败: {str(e)}")
                logger.error(f"异常堆栈: {traceback.format_exc()}")
        else:
            logger.info(f"当前collection '{collection_name}' 不是测试collection，无需删除")

        # 记录初始化时间（删除操作不计入）
        start_time = time.time()

        # 初始化嵌入服务
        try:
            logger.info(f"正在初始化 {config_name} 嵌入服务...")
            embedding_service = EmbeddingService(
                embedding_model_name=config['embedding']['model_name'],
                embedding_model_path=config['embedding']['model_path'],
                rerank_model_name=config['rerank']['model_name'],
                rerank_model_path=config['rerank']['model_path'],
                milvus_collection_name=config['milvus']['collection_name'],
                milvus_connection_args=config['milvus']['connection_args'],
                backend_name=config['backend']['name']
            )
            init_time = time.time() - start_time
            logger.info(f"{config_name} 初始化耗时: {init_time:.4f} 秒")

            # 加载测试文档
            docs = load_test_documents()
            if not docs:
                logger.error(f"未加载到测试文档，无法进行嵌入测试")
                return None

            # 测试文档嵌入 - 使用逐个处理方法
            logger.info(f"开始逐个嵌入 {len(docs)} 个测试文档...")
            start_embed_time = time.time()
            embeddings = []
            for i, doc in enumerate(docs):
                # 使用embed_query方法进行逐个处理
                embedding = embedding_service.embed_query(doc.page_content)
                embeddings.append(embedding)
                if (i+1) % 50 == 0 or (i+1) == len(docs):
                    logger.info(f"已嵌入 {i+1}/{len(docs)} 个文档")
            embed_time = time.time() - start_embed_time
            logger.info(f"{config_name} 嵌入 {len(docs)} 个文档耗时: {embed_time:.4f} 秒")
            logger.info(f"平均每个文档嵌入时间: {embed_time/len(docs):.6f} 秒")

            # 将文档和嵌入向量存储到Milvus
            logger.info(f"将 {len(docs)} 个文档及其嵌入向量存储到Milvus...")
            start_store_time = time.time()
            for i, (doc, embedding) in enumerate(zip(docs, embeddings)):
                # 调用add_document方法将文档和嵌入向量存储到Milvus
                embedding_service.add_document(doc, embedding)
                if (i+1) % 50 == 0 or (i+1) == len(docs):
                    logger.info(f"已存储 {i+1}/{len(docs)} 个文档到Milvus")
            store_time = time.time() - start_store_time
            logger.info(f"存储 {len(docs)} 个文档到Milvus耗时: {store_time:.4f} 秒")
            logger.info(f"平均每个文档存储时间: {store_time/len(docs):.6f} 秒")

            # 加载测试查询
            queries = load_test_queries()
            if not queries:
                logger.error(f"未加载到测试查询，无法进行查询测试")
                return None

            # 测试查询
            start_query_time = time.time()
            logger.info(f"开始处理 {len(queries)} 个查询...")

            # 初始化统计变量
            total_retrieved_docs = 0
            max_retrieved_docs = 0
            min_retrieved_docs = float('inf')

            # 遍历每个查询
            for i, query in enumerate(queries):
                query_start_time = time.time()
                logger.info(f"处理查询 {i+1}/{len(queries)}: {query[:30]}...")

                # 执行检索
                try:
                    # 使用embedding_service的query方法进行检索, 参数包括查询文本、top_k和是否重排序
                    results = embedding_service.query(query, top_k=5, rerank=True)
                    retrieved_count = len(results.get('results', [])) if results.get('success', False) else 0

                    # 更新统计信息
                    total_retrieved_docs += retrieved_count
                    max_retrieved_docs = max(max_retrieved_docs, retrieved_count)
                    min_retrieved_docs = min(min_retrieved_docs, retrieved_count)

                    logger.info(f"查询 {i+1} 检索到 {retrieved_count} 个文档")
                except Exception as e:
                    logger.error(f"查询 {i+1} 处理失败: {str(e)}")

                query_end_time = time.time()
                logger.info(f"查询 {i+1} 耗时: {query_end_time - query_start_time:.4f} 秒")

            query_time = time.time() - start_query_time

            # 计算平均检索文档数
            avg_retrieved_docs = total_retrieved_docs / len(queries) if len(queries) > 0 else 0

            # 打印统计信息
            logger.info(f"{config_name} 处理 {len(queries)} 个查询总耗时: {query_time:.4f} 秒")
            logger.info(f"平均每个查询耗时: {query_time/len(queries):.6f} 秒")
            logger.info(f"总检索文档数: {total_retrieved_docs}")
            logger.info(f"平均每个查询检索文档数: {avg_retrieved_docs:.2f}")
            logger.info(f"最大检索文档数: {max_retrieved_docs}")
            logger.info(f"最小检索文档数: {min_retrieved_docs if min_retrieved_docs != float('inf') else 0}")

            # 清理环境变量
            if 'CONFIG_PATH' in os.environ:
                del os.environ['CONFIG_PATH']
                logger.debug("已清理 CONFIG_PATH 环境变量")

            result = {
                'init_time': init_time,
                'embed_time': embed_time,
                'query_time': query_time,
                'total_time': init_time + embed_time + query_time,
                'doc_count': len(docs),
                'query_count': len(queries)
            }
            logger.info(f"{config_name} 测试完成，结果: {result}")
            return result

        except Exception as e:
            logger.error(f"{config_name} 测试失败: {str(e)}")
            logger.error(f"异常堆栈: {traceback.format_exc()}")
            # 清理环境变量
            if 'CONFIG_PATH' in os.environ:
                del os.environ['CONFIG_PATH']
                logger.debug("已清理 CONFIG_PATH 环境变量")
            return None

    def run_tests(self):
        """运行所有测试"""
        self.results = {}
        logger.info("开始运行模型性能比较测试...")

        # 测试HF配置
        logger.info("准备测试 HuggingFace 配置...")
        hf_result = self.test_config(
            '/Users/apple/Downloads/developments/langchain_rag/config/hf_config.ini', 
            'HuggingFace'
        )
        if hf_result:
            self.results['HuggingFace'] = hf_result
            logger.info(f"HuggingFace 测试结果已添加: {hf_result}")
        else:
            logger.error("HuggingFace 测试失败，未添加结果")

        # 测试FlagModel配置
        logger.info("准备测试 FlagModel 配置...")
        flag_result = self.test_config(
            '/Users/apple/Downloads/developments/langchain_rag/config/flag_config.ini', 
            'FlagModel'
        )
        if flag_result:
            self.results['FlagModel'] = flag_result
            logger.info(f"FlagModel 测试结果已添加: {flag_result}")
        else:
            logger.error("FlagModel 测试失败，未添加结果")

        # 检查结果是否为空
        logger.info(f"当前结果字典: {self.results}")
        if not self.results:
            logger.error("结果字典为空，无法保存到文件")
            return

        # 保存结果
        result_file = os.path.abspath('model_performance_results.json')
        logger.info(f"准备保存结果到文件: {result_file}")
        try:
            # 确保目录存在
            os.makedirs(os.path.dirname(result_file), exist_ok=True)
            with open(result_file, 'w', encoding='utf-8') as f:
                json.dump(self.results, f, ensure_ascii=False, indent=4)
            logger.info(f"结果已成功保存到: {result_file}")
            # 验证文件是否存在
            if os.path.exists(result_file):
                logger.info(f"文件已成功创建，大小: {os.path.getsize(result_file)} 字节")
            else:
                logger.error(f"文件创建失败，路径: {result_file}")
        except Exception as e:
            logger.error(f"保存结果到文件失败: {str(e)}")
            logger.error(f"异常堆栈: {traceback.format_exc()}")

        # 打印比较结果
        logger.info("准备打印性能比较结果...")
        if self.results:
            self.print_comparison()
        else:
            logger.error("结果为空，无法打印比较结果")

    def print_comparison(self):
        """打印性能比较结果"""
        if not self.results:
            logger.info("没有可用的性能测试结果。")
            return

        print("\n===== 模型性能比较结果 =====")
        print(f"{'模型':<15} {'初始化时间(秒)':<15} {'嵌入时间(秒)':<15} {'查询时间(秒)':<15} {'总时间(秒)':<15}")
        print("-" + "="*65 + "")
        for model_name, result in self.results.items():
            print(f"{model_name:<15} {result['init_time']:.4f}        {result['embed_time']:.4f}        {result['query_time']:.4f}        {result['total_time']:.4f}")

        # 计算差异
        if len(self.results) == 2:
            models = list(self.results.keys())
            model1, model2 = models[0], models[1]
            result1, result2 = self.results[model1], self.results[model2]

            print("-" + "="*65 + "")
            print(f"{model1} 相比 {model2} 的性能差异:")
            print(f"初始化时间: {(result1['init_time'] - result2['init_time']) / result2['init_time'] * 100:.2f}%")
            print(f"嵌入时间: {(result1['embed_time'] - result2['embed_time']) / result2['embed_time'] * 100:.2f}%")
            print(f"查询时间: {(result1['query_time'] - result2['query_time']) / result2['query_time'] * 100:.2f}%")
            print(f"总时间: {(result1['total_time'] - result2['total_time']) / result2['total_time'] * 100:.2f}%")

        logger.info("="*75 + "")

# 运行测试
if __name__ == '__main__':
    logger.info("进入主函数")
    try:
        logger.info("创建ModelPerformanceTester实例")
        tester = ModelPerformanceTester()
        logger.info("调用run_tests方法")
        tester.run_tests()
        logger.info("run_tests方法执行完成")
    except Exception as e:
        logger.error(f"主函数执行失败: {str(e)}")
        logger.error(f"异常堆栈: {traceback.format_exc()}")
    logger.info("脚本执行结束")